# -*- coding: utf-8 -*-
# =============================================================================
#         Desc: 格式化MySQL通用日志并统计SQL执行频率
#       Author: GGA
#        Email:
#     HomePage:
#      Version: 1.0.0
#   LastChange: 2020-12-20
#      History:
# =============================================================================
import datetime
import sys
import re
import os
import hashlib
import argparse
import logging
import time
from utils.sqlpf_helper import SQLFingerPrint
from utils.logger_helper import LoggerHelper

default_datetime_format = "%Y-%m-%d %H:%M:%S"

logger = logging.getLogger()


class GeneralLogConfig(object):
    QUERY_START_TAG = " Query\t"
    CONNECT_START_TAG = " Connect\t"
    QUIT_START_TAG = " Quit\t"
    QUERY_TIME_PATTERN = r"^[0-9]{6} [0-9]{2}:[0-9]{1,2}:[0-9]{1,2}"
    QUERY_TIME_PATTERN2 = r"^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}"
    QUERY_TIME_PATTERN3 = r"^[0-9]{6}  [0-9]{1}:[0-9]{1,2}:[0-9]{1,2}"
    CHECK_SQL_LIst = [
        "select @@session.tx_read_only",
        "select 1",
        "set autocommit=1",
        "set autocommit=0",
        "commit"
    ]
    TMP_CACHE_LOG_NUMBER = 10000


class GeneralLogFilter(object):
    def __init__(self, start_time, stop_time, ignore_cheek_sql):
        self.start_time = start_time
        self.stop_time = stop_time
        self.ignore_cheek_sql = ignore_cheek_sql


class GeneralLogHelper(object):

    def __init__(self, source_log_path, log_filter: GeneralLogFilter):
        self.sql_stats_map = {}
        self.source_log_path = source_log_path
        self.log_filter = log_filter
        self.check_time_str = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        self.log_path = os.path.join(os.path.curdir, "logs")
        self.target_log_path = os.path.join(self.log_path, self.check_time_str + "_format.txt")
        self.static_log_path = os.path.join(self.log_path, self.check_time_str + "_stats.txt")
        self.cache_log_items = []

    @classmethod
    def is_check_sql(cls, sql_text: str):
        for tmp_sql in GeneralLogConfig.CHECK_SQL_LIst:
            if sql_text.lower().find(tmp_sql.lower()) >= 0:
                return True
        return False

    def check_env(self):
        if not os.path.exists(self.log_path):
            os.makedirs(self.log_path)
        if not os.path.exists(self.source_log_path):
            logger.info("源文件不存在，请检查")
            return False
        if os.path.exists(self.target_log_path):
            logger.info("目标文件已存在，请检查")
            return False
        if os.path.exists(self.static_log_path):
            logger.info("统计文件已存在，请检查")
            return False
        return True

    def update_sql_map(self, sql_text):
        sql_map = self.sql_stats_map
        dt1 = time.time()
        sql_pf = SQLFingerPrint.get_finger_print(sql=sql_text)
        sql_md5 = hashlib.md5(sql_pf.encode(encoding='UTF-8')).hexdigest()
        dt2 = time.time()
        if sql_md5 in sql_map.keys():
            sql_map[sql_md5]["exec_count"] += 1
        else:
            sql_map[sql_md5] = dict()
            sql_map[sql_md5]["exec_count"] = 1
            sql_map[sql_md5]["sql_md5"] = sql_md5
            sql_map[sql_md5]["finger_print"] = sql_pf
            sql_map[sql_md5]["sample_sql"] = sql_text
        dt3 = time.time()
        if (dt2 - dt1) * 1000 > 10:
            logger.warning("解析SQL耗时超过10ms,SQL:{}".format(sql_text))
        if (dt3 - dt2) * 1000 > 10:
            logger.warning("更新SQL MAP耗时超过10ms,SQL:{}".format(sql_text))

    def get_sql_map_print_info(self, ):
        sql_map = self.sql_stats_map
        print_info = []
        sql_items = sorted(sql_map.values(), key=lambda item: item["exec_count"], reverse=True)
        total_exec_count = 0
        for sql_item in sql_items:
            total_exec_count += sql_item["exec_count"]
        print_info.append("*" * 100 + "")
        print_info.append("开始时间：{}".format(self.log_filter.start_time))
        print_info.append("结束时间：{}".format(self.log_filter.stop_time))
        print_info.append("请求次数：{}".format(total_exec_count))
        if self.log_filter.ignore_cheek_sql == 1:
            print_info.append("检查SQL：忽略")
        else:
            print_info.append("检查SQL：保留")
        for sql_item in sql_items:
            print_info.append("*" * 100 + "")
            print_info.append("指纹SQL:    {0}".format(sql_item["finger_print"]))
            print_info.append("指纹md5:    {0}".format(sql_item["sql_md5"]))
            print_info.append("示例SQL:    {0}".format(sql_item["sample_sql"]))
            print_info.append("执行次数:    {0}".format(sql_item["exec_count"]))
            exec_rate = round(sql_item["exec_count"] * 100.0 / total_exec_count, 2)
            print_info.append("执行占比:    {0}%".format(exec_rate))
        return print_info

    @classmethod
    def append_file_content(cls, file_path, file_content):
        with open(file=file_path, encoding="utf-8", mode="a+") as fw:
            fw.write(file_content)

    def is_keep_sql(self, sql_text, query_time):
        lf = self.log_filter
        if query_time < lf.start_time:
            return False
        if query_time > lf.stop_time:
            return False
        if self.log_filter.ignore_cheek_sql == 0:
            if self.is_check_sql(sql_text=sql_text):
                return False
        return True

    def handler_log_item(self, sql_text, query_time, query_id):
        sql_text = re.sub(
            r"\s+", " ",
            str(sql_text).replace("\r", " ").replace("\n", " ").replace("\t", " ")
        )
        if self.is_keep_sql(sql_text, query_time):
            self.update_sql_map(sql_text)
            self.cache_log_items.append("[{}] [{}] {}".format(query_time, query_id, sql_text))

    def write_cache_logs(self):
        self.append_file_content(
            file_path=self.target_log_path,
            file_content="\n".join(self.cache_log_items)
        )
        self.cache_log_items = []

    def write_status_logs(self):
        print_info = self.get_sql_map_print_info()
        self.append_file_content(
            file_path=self.static_log_path,
            file_content="\n".join(print_info)
        )

    def show_check_result(self):
        logger.info("*" * 40)
        logger.info("统计文件：{}".format(self.static_log_path))
        logger.info("格式化文件：{}".format(self.target_log_path))
        logger.info("处理完成。。。")

    @classmethod
    def try_get_query_time(cls, sql_line):
        try:
            query_time = None
            if re.match(GeneralLogConfig.QUERY_TIME_PATTERN, sql_line):
                query_time_str = re.search(GeneralLogConfig.QUERY_TIME_PATTERN, sql_line).group()
                query_time = datetime.datetime.strptime("20" + query_time_str, "%Y%m%d %H:%M:%S")

            if re.match(GeneralLogConfig.QUERY_TIME_PATTERN2, sql_line):
                query_time_str = re.search(GeneralLogConfig.QUERY_TIME_PATTERN2, sql_line).group()
                query_time = datetime.datetime.strptime(query_time_str, "%Y-%m-%dT%H:%M:%S")

            if re.match(GeneralLogConfig.QUERY_TIME_PATTERN3, sql_line):
                query_time_str = re.search(GeneralLogConfig.QUERY_TIME_PATTERN3, sql_line).group()
                query_time = datetime.datetime.strptime("20" + query_time_str, "%Y%m%d %H:%M:%S")
            return query_time
        except Exception as ex:
            logger.info("处理时间格式出现异常,时间为：{},异常为:{}".format(sql_line, str(ex)))
            exit(-1)

    @classmethod
    def get_sql_tag(cls, sql_line: str):
        sql_tags = [GeneralLogConfig.QUERY_START_TAG, GeneralLogConfig.CONNECT_START_TAG,
                    GeneralLogConfig.QUIT_START_TAG]
        for sql_tag in sql_tags:
            if sql_line.find(sql_tag) > 0:
                return sql_tag
        return None

    @classmethod
    def get_query_id(cls, sql_line):
        sql_tag = cls.get_sql_tag(sql_line=sql_line)
        if sql_tag is None:
            return ""
        query_id_index = str(sql_line).index(sql_tag)
        current_query_id = sql_line[:query_id_index].replace("\t", "#").replace("##", "#").replace("##", "#").strip()
        if current_query_id.find("#") >= 0:
            return current_query_id.split("#")[-1].strip()
        else:
            return current_query_id

    @classmethod
    def get_first_query_text(cls, sql_line):
        sql_tag = cls.get_sql_tag(sql_line=sql_line)
        if sql_tag is None:
            return ""
        sql_index = str(sql_line).index(sql_tag) + len(sql_tag)
        return sql_line[sql_index:]

    def check_logs(self):
        if not self.check_env():
            return
        sql_text = ""
        sql_text_lines = []
        line_num = 0
        new_query_flag = False
        current_query_time = None
        current_query_id = ""
        with open(file=self.source_log_path, encoding="utf-8", mode="r+", errors='ignore') as slh:
            for line in slh:
                line_num += 1
                if line_num % 1000 == 0:
                    logger.info("处理第{}行记录".format(line_num))
                    if len(self.cache_log_items) > GeneralLogConfig.TMP_CACHE_LOG_NUMBER:
                        self.write_cache_logs()
                # 判断此列是否包含时间信息
                tmp_time = self.try_get_query_time(sql_line=line)
                if tmp_time is not None:
                    current_query_time = self.try_get_query_time(sql_line=line)
                # 判断此列是否为SQL开始列
                if line.find(GeneralLogConfig.QUERY_START_TAG) > 0 \
                        or line.find(GeneralLogConfig.CONNECT_START_TAG) > 0 \
                        or line.find(GeneralLogConfig.QUIT_START_TAG) > 0:
                    if current_query_time is None:
                        continue
                    if current_query_id is None:
                        continue
                    if not new_query_flag:
                        sql_text_lines = []
                        new_query_flag = True
                    # 处理上一条SQL
                    sql_text = " ".join(sql_text_lines)
                    self.handler_log_item(
                        sql_text=sql_text,
                        query_time=current_query_time,
                        query_id=current_query_id
                    )
                    # 开始处理新一条SQL
                    sql_text_line = self.get_first_query_text(sql_line=line)
                    sql_text_lines = [sql_text_line]
                    current_query_id = self.get_query_id(sql_line=line)
                else:
                    sql_text_lines.append(line)
        self.write_cache_logs()
        self.write_status_logs()
        self.show_check_result()


def get_parser():
    parser = argparse.ArgumentParser(
        description='格式化MySQL通用日志并汇总SQL执行情况',
        add_help=False)
    connect_setting = parser.add_argument_group('参数信息')
    parser.add_argument(
        '--help',
        dest='help',
        action='store_true',
        help='获取帮助信息',
        default=False
    )
    connect_setting.add_argument(
        '--source_log_path',
        dest='source_log_path',
        type=str,
        default=None,
        help='通用日志文件路径，无默认值，需指定。'
    )
    connect_setting.add_argument(
        '--ignore_cheek_sql',
        dest='ignore_cheek_sql',
        type=int,
        default=1,
        help='是否忽略客户端检查SQL,默认忽略。'
    )
    connect_setting.add_argument(
        '--start_time',
        dest='start_time',
        type=str,
        default="1970-01-01 00:00:00",
        help='日志开始时间,默认为1970-01-01 00:00:00,在该时间点前执行完的慢SQL将被忽略'
    )
    connect_setting.add_argument(
        '--stop_time',
        dest='stop_time',
        type=str,
        default="2199-01-01 00:00:00",
        help='日志结束时间,默认为2199-01-01 00:00:00,在该时间点后执行完的慢SQL将被忽略'
    )
    return parser


def parse_args(command_args):
    need_print_help = False if command_args else True
    parser = get_parser()
    args = parser.parse_args(command_args)
    if args.help or need_print_help:
        parser.print_help()
        sys.exit(1)
    if args.source_log_path is None:
        logger.warning("请输入通用日志文件路径")
        parser.print_help()
        sys.exit(1)
    if args.ignore_cheek_sql not in (1, 0):
        logger.warning("参数ignore_cheek_sql只能赋值0和1")
        parser.print_help()
        sys.exit(1)
    try:
        args.stop_time = datetime.datetime.strptime(str(args.stop_time), default_datetime_format)
        args.start_time = datetime.datetime.strptime(str(args.start_time), default_datetime_format)
    except Exception as ex:
        logger.debug(str(ex))
        logger.warning("请正确的开始时间和结束时间")
        parser.print_help()
        sys.exit(1)
    return args


def main(command_args):
    args = parse_args(command_args)
    lf = GeneralLogFilter(
        start_time=args.start_time,
        stop_time=args.stop_time,
        ignore_cheek_sql=args.ignore_cheek_sql
    )
    lh = GeneralLogHelper(
        source_log_path=args.source_log_path,
        log_filter=lf
    )
    lh.check_logs()


def init_logger():
    LoggerHelper.init_logger(logger_level=logging.INFO)
    global logger
    logger = LoggerHelper.get_logger()


if __name__ == '__main__':
    init_logger()
    main(sys.argv[1:])
